from gensim.models import Word2Vec
from tqdm import tqdm
import re
import jieba
stopwords_path = './data/stopwords.txt'
def tokenizer(text):
    '''分词+去停用词'''
    stopwords = get_stopwords_list()
    sentence_depart = seg_depart(text)
    sentence_depart = move_stopwords(sentence_depart, stopwords)
    return sentence_depart

def get_stopwords_list():
    stopwords = [line.strip() for line in open(stopwords_path, encoding='UTF-8').readlines()]
    return stopwords

def seg_depart(sentence):
    sentence_depart = jieba.lcut(sentence.strip(), cut_all=False, HMM=True)
    return sentence_depart

def remove_digits(input_str):
    punc = u'0123456789.'
    output_str = re.sub(r'[{}]+'.format(punc), '', input_str)
    return output_str

def move_stopwords(sentence_list, stopwords_list):
    out_list = []
    for word in sentence_list:
        if word not in stopwords_list:
            if not remove_digits(word):
                continue
            if word != '\t':
                out_list.append(word)
    return out_list

    
content_clean = []
f = open('./data/article.txt', 'a')
with tqdm(total=2324) as pbar:
    for line in open('./data/news.txt', 'r', encoding='utf8'):
        news = line.split('\'')[1]
        tokens = tokenizer(news)
        content_clean.append(tokens)
        pbar.update(1)
print('>>>')
model = Word2Vec(content_clean, sg=1, vector_size=100, window=5, min_count=2, negative=1, 
                 sample=0.001, workers=4)
print('<<<')
'''
sg=1 是 skip-gram 算法，对低频词敏感；默认 sg=0 为 CBOW 算法。
vector_size 是输出词向量的维数，值太小会导致词映射因为冲突而影响结果，值太大则会耗内存并使算法计算变慢，一般值取为100到200之间。
window 是句子中当前词与目标词之间的最大距离，3表示在目标词前看3-b 个词，后面看 b 个词（b 在0-3之间随机）。
min_count 是对词进行过滤，频率小于 min-count 的单词则会被忽视，默认值为5。
negative 和 sample 可根据训练结果进行微调，sample 表示更高频率的词被随机下采样到所设置的阈值，默认值为 1e-3。
hs=1 表示层级 softmax 将会被使用，默认 hs=0 且 negative 不为0，则负采样将会被选择使用。
'''
model.save('./word2vec')   #保存

model = Word2Vec.load('word2vec')

testwords = ['中印', '台湾', '南海']
for i in range(3):
    res = model.wv.most_similar(testwords[i],topn=10)
    print(testwords[i])
    print(res)